#include "animation.hpp"
#include "json.h"

#include <fstream>

IKinema::Animation::Animation(const std::string& p_filename)
{
	load(p_filename);
}

IKinema::Animation::~Animation()
{}

void IKinema::Animation::load(const std::string& p_filename)
{
	Json::Value result_root{Json::nullValue};
	std::ifstream animation_file{p_filename};
	
	if (!animation_file.is_open())
		throw std::runtime_error{"unable to open animation file"};

	Json::Reader json_reader;
	if (!json_reader.parse(animation_file, result_root))
		throw std::runtime_error{"error parsing file"};

	load_metadata_from(result_root["metadata"]);
	if (m_animated_nodes == 0)
		throw std::runtime_error{"no animated nodes in animation"};

	const Json::Value& anim_data = result_root["animation"];

	if (!anim_data.isArray() || anim_data.size() == 0)
		throw std::runtime_error{"animation must be an non-empty array"};

	for (const auto& frame : anim_data)
		load_frame_from(frame);

	if (m_animation_length == 0)
		throw std::runtime_error{"no animation data read"};
}

const IKinema::vrpn::skeleton_desc_t& IKinema::Animation::get_skeleton_description() const
{
	return m_animation_metadata;
}

const IKinema::vrpn::skeleton_frame_t& IKinema::Animation::get_animation_frame(std::size_t p_requested_frame) const
{
	if (m_animation_length == 0)
		throw std::runtime_error{"animation is empty"};

	return m_anim_frames[p_requested_frame % m_animation_length]; // limit the returned frame to [0, m_animation_length-1]
}

void IKinema::Animation::load_metadata_from(const Json::Value& p_meta)
{
	if (!p_meta.isArray())
		throw std::runtime_error{"metadata must be an array"};

	for (const auto& node : p_meta) {
		vrpn::bone_desc_t bone_desc;
		bone_desc.id = node["id"].asInt();
		bone_desc.parent_id = node["pid"].asInt();
		bone_desc.name = node["name"].asString();
		const auto& rest_transform = node["rest"];
		
		if (!rest_transform.isArray() || rest_transform.size() != 7)
			throw std::runtime_error{"bad animation data"};

		// read translation [X, Y, Z]
		for (std::size_t i = 0; i < 3; ++i)
			bone_desc.rest.translation.data[i] = rest_transform[static_cast<Json::ArrayIndex>(i)].asDouble();

		// read quaternion rotation [X, Y, Z, W]
		for (std::size_t i = 0; i < 4; ++i)
			bone_desc.rest.rotation.data[i] = rest_transform[static_cast<Json::ArrayIndex>(i + 3)].asDouble();

		m_animation_metadata.push_back(std::move(bone_desc));
	}

	m_animated_nodes = m_animation_metadata.size();
}

void IKinema::Animation::load_frame_from(const Json::Value& p_anim_frame)
{
	if (!p_anim_frame.isArray())
		throw std::runtime_error{"animation must be an array"};

	vrpn::skeleton_frame_t new_frame;
	for (const auto& node_transform : p_anim_frame) {
		transform_t local_transform;

		if (!node_transform.isArray() || node_transform.size() != 7)
			throw std::runtime_error{"bad animation data"};

		// read translation [X, Y, Z]
		for (std::size_t i = 0; i < 3; ++i)
			local_transform.translation.data[i] = node_transform[static_cast<Json::ArrayIndex>(i)].asDouble();

		// read quaternion rotation [X, Y, Z, W]
		for (std::size_t i = 0; i < 4; ++i)
			local_transform.rotation.data[i] = node_transform[static_cast<Json::ArrayIndex>(i+3)].asDouble();

		new_frame.push_back(std::move(local_transform));
	}

	if (new_frame.size() != m_animated_nodes)
		throw std::runtime_error{"frame doesn't contain animation for all nodes"};

	m_anim_frames.push_back(std::move(new_frame));
	m_animation_length++;
}

IKinema::Bvh::Bvh (const std::string& bonePrefix, const std::string& p_filename)
{
	m_bonePrefix = bonePrefix;
	load (p_filename);
	BRRegisterFrameDataCallback (this, bvhFrameDataReceived);
	m_socket_ref = nullptr;
	m_tracker = nullptr;
	m_eulerOrder = YXZ;
}

IKinema::Bvh::~Bvh ()
{
	stopService ();
}

const IKinema::vrpn::skeleton_desc_t& IKinema::Bvh::get_skeleton_description () const
{
	return m_animation_metadata;
}

const IKinema::vrpn::skeleton_frame_t& IKinema::Bvh::get_animation_frame () const
{
	return m_animation_frame;
}

const EulerOrder& IKinema::Bvh::eulerOrder () const
{
	return m_eulerOrder;
}

EulerOrder& IKinema::Bvh::eulerOrder ()
{
	return m_eulerOrder;
}

const std::string& IKinema::Bvh::subject () const
{
	return m_subject;
}

std::string& IKinema::Bvh::subject ()
{
	return m_subject;
}

void IKinema::Bvh::startServiceAt (IKinema::vrpn::SkeletonTracker* tracker, int port)
{
	m_tracker = tracker;
	m_socket_ref = BRStartUDPServiceAt (port);
}

void IKinema::Bvh::stopService ()
{
	if (m_socket_ref) {
		BRCloseSocket (m_socket_ref);
		m_socket_ref = nullptr;
	}
}

void IKinema::Bvh::swap_buffer ()
{
}

void IKinema::Bvh::load (const std::string& p_filename)
{
	std::ifstream bvh_file { p_filename, std::ios::binary };

	if (!bvh_file.is_open ())
		throw std::runtime_error{ "unable to open bvh file" };

	long size;
	std::filebuf *pbuf = bvh_file.rdbuf ();
	char* buffer;

	// get file size using buffer's members
	size = pbuf->pubseekoff (0, std::ios::end, std::ios::in);
	pbuf->pubseekpos (0, std::ios::in);

	// allocate memory to contain file data
	buffer = new char[size + 1];
	buffer[size] = 0;

	// get file data  
	pbuf->sgetn (buffer, size);

	C = std::move (std::string (buffer));
	delete[] buffer;

	bvh_file.close ();

	parse ();

	build_tpose ();
}

void IKinema::Bvh::parse ()
{
	prolog ();
	matchString ("HIERARCHY");
	joints ();
}

void IKinema::Bvh::joints () {
	std::string joint_begin = name ();
	if (joint_begin == "ROOT" ||
		joint_begin == "JOINT") {

		vrpn::bone_desc_t bone_desc;
		bone_desc.name = m_bonePrefix+name ();
		bone_desc.id = BONE_ID++;

		if (STACK.empty ()) {
			bone_desc.parent_id = -1;
		}
		else {
			bone_desc.parent_id = STACK.back ().id;
		}

		match ('{');
		auto& T = bone_desc.rest.translation.data;
		float x, y, z;
		offset (x, y, z);
		T[0] = x;
		T[1] = y;
		T[2] = z;

		auto& Q = bone_desc.rest.rotation.data;
		Q[0] = Q[1] = Q[2] = 0;
		Q[3] = 1;	// w-component

		STACK.push_back (bone_desc);
		m_animation_metadata.push_back (bone_desc);

		skipToChar ('\n');
		skipWhite ();

		// parse sub-node
		int P_backup = P;
		joint_begin = name ();
		if (joint_begin == "JOINT") {
			P = P_backup;
			joints ();
		}
		else if (joint_begin == "End") {
			skipToChar ('}');
			match ('}');
		}

		match ('}');
		STACK.pop_back ();

		// parse sibling-node
		if (!isEOF () && isalpha (C[P])) {
			joints ();
		}
	}
}

void IKinema::Bvh::prolog ()
{
	BONE_ID = 0;
	L = C.size ();
	P = -1;
	STACK.clear ();
	getChar ();
	skipWhite ();
	m_animation_metadata.clear ();
}

void IKinema::Bvh::match (char ch)
{
	if (!isEOF () && ch == C[P]) {
		getChar ();
		skipWhite ();
	}
}

void IKinema::Bvh::matchString (const std::string& str)
{
	if (L - P >= str.size ()) {
		for (int i = 0; i < str.size (); ++i) {
			if (C[P] == str[i]) {
				getChar ();
			}
		}
		skipWhite ();
	}
}

void IKinema::Bvh::skipWhite ()
{
	while (!isEOF () && isspace (C[P])) {
		getChar ();
	}
}

void IKinema::Bvh::skipToChar (char ch)
{
	while (!isEOF () && (C[P] != ch)) {
		getChar ();
	}
}

void IKinema::Bvh::skipToString (const std::string& str)
{
	P = C.find_first_of (str, P);
}

void IKinema::Bvh::getChar ()
{
	++P;
}

bool IKinema::Bvh::isEOF ()
{
	return P >= L;
}

std::string IKinema::Bvh::name () {
	std::string s_name;
	while (!isEOF () && (isalnum(C[P]) || C[P] == '_')) {
		s_name += C[P];
		getChar ();
	}

	skipWhite ();
	return s_name;
}

float IKinema::Bvh::number () {
	float num = 0.0f;

	std::string s_num = "";
	if (C[P] == '-') {
		s_num += C[P];
		getChar ();
		skipWhite ();
	}

	while (!isEOF () && (isdigit (C[P]) || C[P] == '.')) {
		s_num += C[P];
		getChar ();
	}

	num = atof (s_num.c_str());

	skipWhite ();
	return num;
}

void IKinema::Bvh::offset (float& x, float& y, float& z)
{
	matchString ("OFFSET");
	x = number ();
	y = number ();
	z = number ();
}

void IKinema::Bvh::init_global_trans ()
{
	m_animation_frame.resize (m_animation_metadata.size ());
	m_filled.swap (std::vector<bool> (m_animation_metadata.size (), false));
}

void IKinema::Bvh::calc_local_quat (int i)
{
	int id = i;

	if (!m_filled[id]) {

		if (id == 0) {
			m_filled[id] = true;
			if (m_withReference) {
				m_bone_transforms[id].t += m_reference_transform.t;
				m_bone_transforms[id].q = m_reference_transform.q * m_bone_transforms[id].q;
			}
			return;
		}
		else {
			int parent_id = m_animation_metadata[id].parent_id;

			calc_local_quat (parent_id);
			m_bone_transforms[id].q = m_bone_transforms[parent_id].q * m_bone_transforms[id].q;

			// T-Pose bone local translation
			auto& childLocalT = m_animation_metadata[id].rest.translation.data;
			auto& parentLocalT = m_animation_metadata[parent_id].rest.translation.data;

			Eigen::Vector3f parentToChildVector (
				childLocalT[0] - parentLocalT[0],
				childLocalT[1] - parentLocalT[1],
				childLocalT[2] - parentLocalT[2]
				);
			parentToChildVector = m_bone_transforms[parent_id].q._transformVector (parentToChildVector);

			m_bone_transforms[id].t = m_bone_transforms[parent_id].t + parentToChildVector;

			m_filled[id] = true;
			return;
		}
	}
}

void IKinema::Bvh::build_tpose () {
	m_animation_frame.resize (m_animation_metadata.size ());

	for (int i = 0; i != m_animation_metadata.size (); i++) {

		const auto& node_transform = m_animation_metadata[i];

		transform_t local_transform;

		// read translation [X, Y, Z]
		for (std::size_t i = 0; i < 3; ++i) {
			local_transform.translation.data[i] = node_transform.rest.translation.data[i];
		}

		// read quaternion rotation [X, Y, Z, W]
		for (std::size_t i = 0; i < 4; ++i) {
			local_transform.rotation.data[i] = node_transform.rest.rotation.data[i];
		}

		m_animation_frame[i] = std::move (local_transform);
	}
}

void __stdcall IKinema::Bvh::bvhFrameDataReceived (void* customedObj, SOCKET_REF sender, BvhDataHeader* header, float* data) {

	IKinema::Bvh* bvh = static_cast<IKinema::Bvh*>(customedObj);

	auto& head = *header;
	auto& body = data;
	int dataIndex;

	if (std::string ((char*)head.AvatarName) == bvh->m_subject) {

		int bone_count = 0;
		if (head.WithDisp) {
			if (head.WithReference) {
				bone_count = (head.DataCount - 6) / 6;
			}
			else {
				bone_count = head.DataCount / 6;
			}
		}
		else {
			if (head.WithReference) {
				bone_count = ((head.DataCount - 6) - 3) / 3;
			}
			else {
				bone_count = (head.DataCount - 3) / 3;
			}
		}

		std::vector<stTransform>& ts = bvh->m_bone_transforms;
		ts.clear ();
		ts.reserve (bone_count);

		stTransform& refTrans = bvh->m_reference_transform;
		if (head.WithReference) {
			bvh->m_withReference = true;

			// write reference data
			float dispX = body[0];	// x
			float dispY = body[1];	// y
			float dispZ = body[2];	// z
			float ang0 = body[3];
			float ang1 = body[4];
			float ang2 = body[5];

			refTrans.t[0] = dispX;
			refTrans.t[1] = dispY;
			refTrans.t[2] = dispZ;

			refTrans.q = sortedEuler2Quat (a2r (ang0), a2r (ang1), a2r (ang2), bvh->eulerOrder());
		}
		else {
			bvh->m_withReference = false;
		}

		if (head.WithDisp)
		{
			for (int curSel = 0; curSel < bone_count; curSel++) {
				dataIndex = curSel * 6;
				if (head.WithReference)
				{
					dataIndex += 6;
				}

				float dispX = body[dataIndex + 0];
				float dispY = body[dataIndex + 1];
				float dispZ = body[dataIndex + 2];
				float ang0 = body[dataIndex + 3];
				float ang1 = body[dataIndex + 4];
				float ang2 = body[dataIndex + 5];

				/*if (curSel == 0) {
					static char strbuf[128];
					sprintf (strbuf, "x:%f y:%f z:%f\n", dispX, dispY, dispZ);
					OutputDebugString (strbuf);
				}*/

				stTransform local_trans;

				local_trans.t[0] = dispX;
				local_trans.t[1] = dispY;
				local_trans.t[2] = dispZ;

				local_trans.q = sortedEuler2Quat (a2r (ang0), a2r (ang1), a2r (ang2), bvh->eulerOrder ());

				ts.push_back (std::move (local_trans));
			}
		}
		else	// never be here. The motion data must contain displacement.
		{
			//throw std::runtime_error ("The motion data MUST contain displacement.");

			// write hip's data
			{
				dataIndex = 0;
				if (head.WithReference)
				{
					dataIndex += 6;
				}

				// hip's displacement
				float dispX = body[dataIndex + 0];
				float dispY = body[dataIndex + 1];
				float dispZ = body[dataIndex + 2];
				// hip's angle
				float ang0 = body[dataIndex + 3];
				float ang1 = body[dataIndex + 4];
				float ang2 = body[dataIndex + 5];

				stTransform local_trans;

				local_trans.t[0] = dispX;
				local_trans.t[1] = dispY;
				local_trans.t[2] = dispZ;

				local_trans.q = sortedEuler2Quat (a2r (ang0), a2r (ang1), a2r (ang2), bvh->eulerOrder ());

				ts.push_back (std::move (local_trans));
			}

			for (int curSel = 1; curSel < bone_count; curSel++) {
				dataIndex = 3 + curSel * 3;
				if (head.WithReference)
				{
					dataIndex += 6;
				}

				// show angle
				float ang0 = body[dataIndex + 0];
				float ang1 = body[dataIndex + 1];
				float ang2 = body[dataIndex + 2];

				stTransform local_trans;
				local_trans.q = sortedEuler2Quat (a2r (ang0), a2r (ang1), a2r (ang2), bvh->eulerOrder ());

				ts.push_back (std::move (local_trans));
			}
		}

		// rearrange frame to generic mocap data format
		bvh->init_global_trans ();
		for (int i = 0; i != bvh->m_bone_transforms.size (); i++) {
			//bvh->calc_local_quat (i);

			const auto& node_transform = bvh->m_bone_transforms[i];

			transform_t local_transform;

			// read translation [X, Y, Z]
			for (std::size_t i = 0; i < 3; ++i) {
				local_transform.translation.data[i] = node_transform.t[i];
			}

			// read quaternion rotation [X, Y, Z, W]
			for (std::size_t i = 0; i < 4; ++i) {
				local_transform.rotation.data[i] = node_transform.q.coeffs ().coeff (i);
			}

			bvh->m_animation_frame[i] = std::move (local_transform);
		}

		if (bvh->m_tracker) {
			bvh->m_tracker->send_frame (bvh->m_animation_frame);
		}
	}
}

